# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import getpass
import logging
import os
import subprocess
import sys
import tempfile

from absl import app, flags
import jinja2

from coin.base.datetime_util import iterate_date
from coin.support.feed_tool.feed_stats.app.feed_stats_runner import EXCHANGE_LIST


class NoJobError(Exception):
  pass


def generate_slurm_script(start_date,
                          end_date,
                          db_config,
                          git_commit,
                          summary_mailto,
                          job_submit_period=None,
                          outdir_prefix=None,
                          user=None,
                          grid=None,
                          venv=None,
                          flow_root_dir=None,
                          tag=None):
  job_submit_period = 0.01 if job_submit_period is None else job_submit_period
  grid = grid or 'iosg'
  outdir_prefix = (
      outdir_prefix
      or '/remote/{grid}/data/etc/general_operations/deploy/bot-hft/hft.feed_stats/coin/scratch'.
      format(grid=grid))

  out_dir = '%s/feed_stats-${SLURM_JOB_ID}' % outdir_prefix

  start_date = datetime.datetime.strptime(start_date, '%Y%m%d')
  end_date = datetime.datetime.strptime(end_date, '%Y%m%d')
  trading_dates = [td for td in iterate_date(start_date.date(), end_date.date())]
  trading_dates = [td.strftime('%Y%m%d') for td in trading_dates]

  machines = [
      "feed-01.ap-northeast-2.aws",
      "feed-04.ap-northeast-1.aws",
  ]

  if tag:
    mail_subject = "[coin-feed-stats] ${SLURM_JOB_ID} summary - %s" % tag
  else:
    mail_subject = "[coin-feed-stats] ${SLURM_JOB_ID} summary"

  template_loader = jinja2.FileSystemLoader(searchpath="./")
  template_env = jinja2.Environment(loader=template_loader)

  template = template_env.get_template(
      'coin/support/feed_tool/feed_stats/app/feed_stats_slurm/slurm_script.tmpl')
  return template.render(user=user,
                         venv=venv,
                         out_dir=out_dir,
                         job_submit_period=job_submit_period,
                         trading_dates=trading_dates,
                         machines=machines,
                         exchange_list=EXCHANGE_LIST,
                         db_config=db_config,
                         git_commit=git_commit,
                         mail_subject=mail_subject,
                         summary_mailto=summary_mailto)


class SBatchError(Exception):
  pass


def run_sbatch(script_filepath, sbatch_args=None, verbose=False):
  sbatch_args = sbatch_args or []
  if isinstance(sbatch_args, str):
    sbatch_args = [token for token in sbatch_args.split(' ') if token]

  sbatch_args = ['--parsable'] + sbatch_args
  if sbatch_args:
    logging.debug('sbatch_args: %s', ' '.join(sbatch_args))

  completed = subprocess.run(['sbatch', *sbatch_args, script_filepath], stdout=subprocess.PIPE)
  if completed.returncode != 0:
    raise SBatchError(completed)

  sbatch_output = completed.stdout.decode('utf-8')
  job_id = int(sbatch_output.strip())
  if verbose:
    sys.stderr.write('Job ID: %d\n' % job_id)
  return job_id


def run_feed_stats_on_slurm(start_date,
                            end_date,
                            db_config,
                            git_commit,
                            job_submit_period=None,
                            outdir_prefix=None,
                            user=None,
                            grid=None,
                            venv=None,
                            force_use_venv=True,
                            summary_mailto='{user}@prestolabs.io',
                            tag=None,
                            flow_root_dir=None,
                            dry_run=False,
                            sbatch_args=None,
                            notify_user=None,
                            verbose=False):
  if force_use_venv and venv is None:
    user = user or getpass.getuser()
    grid = grid or 'iosg'
    venv = os.environ.get(
        'COIN_SLURM_VENV',
        '/remote/iosg/data/etc/general_operations/deploy/bot-hft/hft.feed_stats/coin/venv'.format(
            grid=grid))

  if summary_mailto is not None:
    if '{user}' in summary_mailto:
      user = notify_user or 'leon'
      summary_mailto = summary_mailto.format(user=user)
    if verbose:
      sys.stderr.write('Feed stats result mail will be sent to %s.\n' % summary_mailto)

  try:
    script = generate_slurm_script(start_date,
                                   end_date,
                                   db_config,
                                   git_commit,
                                   summary_mailto,
                                   job_submit_period=job_submit_period,
                                   outdir_prefix=outdir_prefix,
                                   user=user,
                                   grid=grid,
                                   venv=venv,
                                   tag=tag,
                                   flow_root_dir=flow_root_dir)
  except NoJobError:
    if verbose:
      sys.stderr.write('No job to run.\n')
    raise

  if dry_run:
    if verbose:
      print(script)
    return

  with tempfile.NamedTemporaryFile(mode='w') as script_file:
    script_file.write(script)
    script_file.flush()

    try:
      job_id = run_sbatch(script_file.name, sbatch_args, verbose=verbose)
      if verbose:
        print(job_id)
    except SBatchError:
      if verbose:
        sys.stderr.write('Error while running sbatch. Abort.\n')
      raise

  return job_id


def main(argv):
  FLAGS = flags.FLAGS

  venv = FLAGS.venv or os.environ.get('COIN_SLURM_VENV', None)
  try:
    run_feed_stats_on_slurm(start_date=FLAGS.start_date,
                            end_date=FLAGS.end_date,
                            db_config=FLAGS.db_config,
                            git_commit=FLAGS.git_commit,
                            outdir_prefix=FLAGS.outdir,
                            grid=FLAGS.grid,
                            venv=venv,
                            force_use_venv=False,
                            summary_mailto=FLAGS.mailto,
                            tag=FLAGS.tag,
                            flow_root_dir=FLAGS.flow_root_dir,
                            dry_run=FLAGS.script_only,
                            sbatch_args=FLAGS.sbatch_args,
                            notify_user=FLAGS.notify_user,
                            verbose=True)
  except NoJobError:
    return 0
  except SBatchError:
    return 1


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')
  logging.getLogger('coin.exchange.util.feed_checker').setLevel(logging.WARNING)

  flags.DEFINE_string('start_date', None, 'yyyymmdd')
  flags.DEFINE_string('end_date', None, 'yyyymmdd. Feed at end_date is not checked.')
  flags.DEFINE_string('db_config', None, 'Database config file.')
  flags.DEFINE_string('git_commit', 'HEAD', 'Git commit sha or HEAD.')
  flags.DEFINE_string('outdir',
                      None,
                      'Output directory. Default is /remote/<grid>/home/<user>/scratch')
  flags.DEFINE_string('grid', 'iosg', 'Specify slurm grid you will use.')
  flags.DEFINE_string(
      'venv',
      None,
      'If specified, use the given virtual env. '
      'User can use venv by setting COIN_SLURM_VENV environment variable.')
  flags.DEFINE_string('mailto',
                      '{user}@prestolabs.io',
                      'If specified, it will send a mail after sim finishes.')
  flags.DEFINE_string('tag', None, 'If specified, tag will be in a result mail subject')
  flags.DEFINE_bool('script_only',
                    False,
                    'If true, it will print sbatch script and not run sbatch.')
  flags.DEFINE_string('flow_root_dir',
                      None,
                      'If specified, use the specified directory as flow roor directory.')
  flags.DEFINE_string('sbatch_args', '-n 20', 'sbatch arguments')
  flags.DEFINE_string('notify_user', 'leon', 'Email notification receiver.')
  app.run(main)
